import os
import random
import shutil

def split_dataset(input_dir, output_dir, train_ratio=0.8):
    """
    将数据集按指定比例划分为训练集和测试集
    :param input_dir: 原始数据集文件夹路径
    :param output_dir: 输出数据集文件夹路径
    :param train_ratio: 训练集所占比例（默认0.8）
    """
    # 创建输出文件夹
    train_dir = os.path.join(output_dir, 'trainA')
    test_dir = os.path.join(output_dir, 'testA')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 获取所有图片文件名
    all_images = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
    random.shuffle(all_images)  # 随机打乱图片顺序

    # 计算训练集和测试集的分界点
    split_idx = int(len(all_images) * train_ratio)
    train_images = all_images[:split_idx]
    test_images = all_images[split_idx:]

    # 将图片复制到对应的文件夹中
    for img_name in train_images:
        shutil.copy(os.path.join(input_dir, img_name), os.path.join(train_dir, img_name))

    for img_name in test_images:
        shutil.copy(os.path.join(input_dir, img_name), os.path.join(test_dir, img_name))

    print(f"数据集划分完成！")
    print(f"训练集：{len(train_images)} 张图片，保存到 {train_dir}")
    print(f"测试集：{len(test_images)} 张图片，保存到 {test_dir}")

# 示例使用
input_dir = "rawA_enhance"      # 原始数据集文件夹路径
output_dir = r"D:\develop\PythonCode\python基础\附_项目实战\十四_给CycleGan网络模型制作数据集\blackpoint_mini"      # 输出数据集文件夹路径
split_dataset(input_dir, output_dir, train_ratio=0.8)
